import os
import time
import datetime
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import requests
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import matplotlib.pyplot as plt

scaler = MinMaxScaler()
# 检查是否有可用的 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
torch.set_num_threads(4)
class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SimpleNet, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)
    
    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        output = self.fc(lstm_out[:, -1, :])  # 取最后一个时间步的输出
        return output

    def preprocess_data(self, data):
        """
        预处理股票数据，生成特征和目标值
        :param data: 包含股票数据的列表
        :return: 特征张量和目标张量
        """
        features = []
        targets = []
        for i in range(len(data) - 5):  # 确保有足够的数据来形成6天的序列
            input_data = []
            for j in range(i, i + 5):  # 前5天作为特征
                day_data = data[j]
                # if 'ma_price5' not in day_data:
                #     print(f"Skipping day {j} due to missing ma_price5",day_data)
                #     continue
                input_data.append([
                    float(day_data['open']),
                    float(day_data['high']),
                    float(day_data['low']),
                    float(day_data['close']),
                    float(day_data['volume']),
                    # day_data['ma_price5'],
                    # day_data['ma_volume5']
                ])
            features.append(input_data)
            # 第6天的收盘价作为目标
            target_data = data[i + 5]
            targets.append(float(target_data['close']))

        # 转换为张量
        # print('features',features)
        features = torch.tensor(features, dtype=torch.float32).view(-1, 5 * 5)  # 5天 * 5特征
        targets = torch.tensor(targets, dtype=torch.float32).view(-1, 1)  # 目标值为1维
        print('targets',targets)
        # features = torch.tensor(features, dtype=torch.float32)
        return features, targets

    def train_model(self, train_loader, num_epochs=10, learning_rate=0.001):
        """
        训练模型的方法
        :param train_loader: DataLoader，训练数据加载器
        :param num_epochs: int，训练轮数
        :param learning_rate: float，学习率
        """
        # 定义损失函数和优化器
        criterion = nn.MSELoss()  # 回归问题使用均方误差损失
        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)

        self.train()  # 设置模型为训练模式
        for epoch in range(num_epochs):
            running_loss = 0.0
            for inputs, targets in train_loader:
                # inputs.to(device), targets.to(device)
                # 前向传播
                outputs = self(inputs)
                loss = criterion(outputs, targets)

                # 反向传播和优化
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                running_loss += loss.item() * inputs.size(0)

            epoch_loss = running_loss / len(train_loader.dataset)
            print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss:.4f}")
    def predict(self, input_data):
        """
        预测方法
        :param input_data: 包含五天数据的列表，每个元素是一个字典
        :return: 预测结果（第6天的收盘价）
        """
        # 将输入数据转换为张量
        input_features = []
        for day_data in input_data:
            input_features.append([
                float(day_data['open']),
                float(day_data['high']),
                float(day_data['low']),
                float(day_data['close']),
                float(day_data['volume']),
                # day_data['ma_price5'],
                # day_data['ma_volume5']
            ])
        # print('randn',torch.randn(1, 25))
        input_features = torch.tensor(input_features, dtype=torch.float32).view(1, -1)  # 形状为 (1, 35)
        print('input_features',input_features)

        # input_features = input_features.to(device)  # 将数据移动到设备
        # 设置模型为评估模式
        self.eval()
        with torch.no_grad():  # 关闭梯度计算
            # print(self)
            prediction = self(input_features)
        return prediction.item()  # 返回预测值

def getStockData():
    url = "https://api.mairui.club/hslt/list/b997d4403688d5e66a"  # 替换为实际的licence证书
    response = requests.get(url)
    data = response.json()

    # for stock in data[0:3]:
    #     print(f"股票代码：{stock['dm']}, 股票名称：{stock['mc']}, 交易所：{stock['jys']}")
    return data
    pass
def getHistoryData(code):
    url = "http://money.finance.sina.com.cn/quotes_service/api/json_v2.php/CN_MarketData.getKLineData"
    params = {
        "symbol": code,
        "scale": "240",  # 时间周期，如5分钟、日线等
        "ma": "5",
        "datalen": "1024"
    }
    response = requests.get(url, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        print('error',code,response.status_code)
    # totalPrice = 0
    # for item in data:
        # totalPrice += float(item['close'])
        # print(f"日期：{item['d']}, 开盘价：{item['o']}, 最高价：{item['h']}, 最低价：{item['l']}, 收盘价：{item['c']}")
        # print(data)
    # totalPrice /= len(data)
    # print('totalPrice',totalPrice)

# 数据预处理
scaler = MinMaxScaler()
features = scaler.fit_transform(features)


# 创建模型实例
model = LSTMNet(input_dim=5, hidden_dim=128, output_dim=1)
model.to(device)

# 训练模型
model.train_model(train_loader, num_epochs=1000, learning_rate=0.001)

# # 创建模型实例
# model = SimpleNet()

# 加载模型权重
if os.path.exists("model_weights.pth"):
    state_dict = torch.load("model_weights.pth")
    # print('keys',state_dict.keys())
    # print('fc1',state_dict['fc1.weight'])
    model.load_state_dict(state_dict)
    model.to(device)  # 确保模型移动到正确的设备
    print("Model weights loaded.")
def dotrain():

    stocks = getStockData()
    stocks = stocks[40:60]
    # stocks = [{'dm': '605100', 'mc': '华丰股份', 'jys': 'sh'}]
    # stocks = [{'dm': '601933', 'mc': '永辉超市', 'jys': 'sh'}]
    total = len(stocks)
    for index,stock in enumerate(stocks):
        data = getHistoryData(f"{stock['jys']}{stock['dm']}")
        if data is None:
            continue
        # data = data[0:5]
        # print(data)
        # 预处理数据
        features, targets = model.preprocess_data(data)
        # print('data',len(data))
        # print('features',(features[0]))
        # print('targets',(targets[0]))
        # 封装为 DataLoader
        dataset = TensorDataset(features, targets)
        train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

        time.sleep(1)
        # 调用训练方法
        model.train_model(train_loader, num_epochs=400, learning_rate=0.00001)
        print(f"训练进度:",index/total*100,"%")

    # 假设模型已经训练完成
    torch.save(model.state_dict(), "model_weights.pth")
    print("Model weights saved.")
def dopredict():
    input_data = getHistoryData('sh605100')
    if input_data is None:
        print('error input_data None')
        return
    offset = 0
    data1 = input_data[-6-offset:-1-offset]
    data2 = input_data[-6-offset:-1-offset]
    print('input_data',data1)
    # for item in data1:
    #     print(f"日期：{item['day']}, 收盘价：{item['close']}")
    # predicted_price = model.predict(data1)
    # print(f"Predicted closing price for the next day: {predicted_price:.2f}")

    # 预测下一天的价格
    last_date_str = data1[-1]['day']
    predicted_price = model.predict(data1)
    print(f"预测 {last_date_str} 的收盘价为 : {predicted_price}")
    # print(f"Predicted closing price for the next day: {predicted_price:.2f}")
    """
    for _ in range(100):  # 重复预测5天
        # print(f"Processing day {_ + 1}:")
        # for item in data1:
        #     print(f"日期：{item['day']}, 收盘价：{item['close']}")

        # 获取最后一个数据的日期
        last_date_str = data1[-1]['day']
        last_date = datetime.datetime.strptime(last_date_str, "%Y-%m-%d").date()
        # 预测下一天的价格
        predicted_price = model.predict(data1)
        # print(f"Predicted closing price for the next day: {predicted_price:.2f}")
        if(predicted_price > float(data1[-1]['close'])):
            print(f"预测 {last_date_str} 的收盘价为 : {predicted_price:.2f},涨")
        # 将预测结果追加到数组末尾
        # 计算下一个真实日期
        next_date = last_date + datetime.timedelta(days=1)
        next_date_str = next_date.strftime("%Y-%m-%d")
        new_day_data = {
            'day': next_date_str,  # 假设日期为'next_day'，可以根据需要修改
            'close': predicted_price
        }
        data1.append(new_day_data)
        data2.append(new_day_data)
        # 移除数组的第一个元素
        data1.pop(0)
    # 初始化用于绘图的数据
    dates = [item['day'] for item in data2]  # 日期
    closing_prices = [float(item['close']) for item in data2]  # 收盘价
    # 绘制收盘价图表
    plt.figure(figsize=(12, 6))
    plt.plot(dates, closing_prices, marker='o', linestyle='-', color='b', label='Closing Prices')
    plt.title('Stock Closing Prices and Predictions')
    plt.xlabel('Date')
    plt.ylabel('Closing Price')
    plt.xticks(rotation=45)  # 旋转日期标签，避免重叠
    plt.legend()
    plt.tight_layout()
    plt.show()
        # print(f"Updated input_data: {data1}")
        
    """
dotrain()
dopredict()